import pandas as pd


class Processor:
    def __init__(self, file_path, test=False):
        self.file_path = file_path
        self.is_test = test

    def process(self):
        if self.is_test:
            data = pd.read_table(self.file_path, header=None)
            data.columns = ["s1", "s2"]
            sentence1 = data["s1"].tolist()
            sentence2 = data["s2"].tolist()
            # labels = []
            return sentence1, sentence2, None

        data = pd.read_table(self.file_path, header=None).fillna("0")
        data.columns = ["s1", "s2", "label"]
        sentence1 = data["s1"].tolist()
        sentence2 = data["s2"].tolist()
        labels = data["label"].tolist()

        return sentence1, sentence2, labels


if __name__ == '__main__':
    path = '../gaiic_track3_round1_train_20210228.tsv'
    print(Processor(path).process())
